"""
基金数据采集器模块

实现基金、ETF等产品的数据采集功能
"""

import akshare as ak
import datetime
import logging
import pandas as pd
from typing import List, Dict
from sqlalchemy.orm import Session
from sqlalchemy.exc import IntegrityError

from modules.models.base import get_db
from modules.models.stock import SecurityInfo, DailyQuote, CollectionTask
from config.settings import DATA_COLLECTION

logger = logging.getLogger("fund_collector")

class FundCollector:
    """基金数据采集类"""
    
    def __init__(self):
        self.retry_limit = DATA_COLLECTION["retry_limit"]
        self.retry_delay = DATA_COLLECTION["retry_delay"]
        self.db = next(get_db())
    
    def __del__(self):
        if hasattr(self, 'db') and self.db:
            self.db.close()

    def collect_etf_daily(self, symbol: str):
        """采集ETF日行情数据"""
        try:
            df = ak.fund_em_etf_info(symbol=symbol)
            if df.empty:
                return 0
                
            # 数据清洗
            df = df.rename(columns={
                '净值日期': 'trade_date',
                '单位净值': 'nav',
                '累计净值': 'accumulated_nav',
                '日增长率': 'daily_growth'
            })
            df['trade_date'] = pd.to_datetime(df['trade_date'])
            
            # 保存到数据库
            # 由于我们没有ETFDailyQuote模型，这里使用DailyQuote模型代替
            existing = self.db.query(DailyQuote.date).filter(
                DailyQuote.symbol == symbol
            ).all()
            existing_dates = {x[0] for x in existing}
            
            new_records = []
            for _, row in df.iterrows():
                if row['trade_date'].date() not in existing_dates:
                    new_records.append(DailyQuote(
                        symbol=symbol,
                        date=row['trade_date'],
                        open=row['nav'],
                        high=row['nav'],
                        low=row['nav'],
                        close=row['nav'],
                        volume=0,
                        amount=0
                    ))
            
            if new_records:
                self.db.add_all(new_records)
                self.db.commit()
            
            return len(new_records)
            
        except Exception as e:
            self.db.rollback()
            logger.error(f"ETF数据采集失败: {str(e)}")
            return 0

    def collect_fund_basic_info(self):
        """采集全市场基金基本信息"""
        try:
            df = ak.fund_em_fund_info()
            # 数据清洗和处理逻辑...
            # 保存到SecurityInfo模型
            return True
        except Exception as e:
            logger.error(f"基金信息采集失败: {str(e)}")
            return False

    def collect_fund_holdings(self, fund_code: str):
        """采集基金持仓数据"""
        try:
            df = ak.fund_em_portfolio_hold(fund=fund_code)
            # 解析持仓数据...
            # 保存到CollectionTask模型
            return True
        except Exception as e:
            logger.error(f"持仓数据采集失败: {str(e)}")
            return False